#include <cstring>
#include <string>
#include <cstdio>
#include <iostream>
#include <TFile.h>
#include <TTree.h>
#include <vector>
#include <TCanvas.h>
#include <TH1F.h>
#include <TLegend.h>
#include <TStyle.h>
#include <sstream>

/*
  Script to create histograms from .lhe and .lhco event files
 */

float getPz(float eta, float pT) {
  float theta = 2*std::atan(std::exp(-eta));
  if(eta>=0) {
    theta+=atan(0);
  }
  return pT/tan(theta);
}

void histogramCreationCombined(char** fileNames, char** fileTypes, char** descriptions)
{
  std::cout << "Histogram creation started!" << std::endl;
  TStyle* myStyle = new TStyle("Plain","Style1");

  const int NFILES = 2;
  int types [NFILES] = {std::atoi(fileTypes[0]),std::atoi(fileTypes[1])};

  int NFILESLHE = 0;
  int NFILESLHCO = 0;
 
  for(int i = 0; i < NFILES; i++) {
    std::cout << "Reading from file: " << fileNames[i] << " type: " << types[i] << " description: " << descriptions[i]<< std::endl;
    if(types[i]==0) {
      NFILESLHE++;
    }
    else if(types[i]==1) {
      NFILESLHCO++;
    }
  }
  if(NFILESLHCO+NFILESLHE!=NFILES) {
    std::cout << "ERROR: NFILESLHCO+NFILESLHE!=NFILES" << std::endl;
  }

  TTree **t = new TTree*[NFILESLHE];
  TTree **t2 = new TTree*[NFILESLHCO];
  int iNFILESLHE = 0;
  int iNFILESLHCO = 0;
  for(int i = 0; i<NFILES;i++) {
    if(types[i]==0) {
      TFile *f = TFile::Open(fileNames[i],"READ");
      f->GetObject("tree",t[iNFILESLHE]);
      iNFILESLHE++;
    }
    else if(types[i]==1) {
      TFile *f2 = TFile::Open(fileNames[i],"READ");
      f2->GetObject("tree",t2[iNFILESLHCO]);
      iNFILESLHCO++;
    }
  }

  std::vector<int> *vdirection = 0;
  std::vector<float> *veta = 0;
  std::vector<float> *vphi = 0;
  std::vector<float> *vM = 0;
  std::vector<int> *vpdgId = 0;
  std::vector<float> *vpT = 0;
  std::vector<float> *vp = 0;

  std::vector<int> *typ = 0;
  std::vector<float> *eta = 0;
  std::vector<float> *phi = 0;
  std::vector<float> *pt = 0;
  std::vector<float> *jmas = 0;
  std::vector<float> *ntrk = 0;
  std::vector<float> *btag = 0;
  std::vector<float> *hadEm = 0;

  TCanvas *c1 = new TCanvas("c1","MadGraph analysis",200,10,700,500);

  TBranch **bvdirection = new TBranch*[NFILESLHE];
  TBranch **bveta = new TBranch*[NFILESLHE];
  TBranch **bvphi = new TBranch*[NFILESLHE];
  TBranch **bvM = new TBranch*[NFILESLHE];
  TBranch **bvpdgId = new TBranch*[NFILESLHE];
  TBranch **bvpT = new TBranch*[NFILESLHE];
  TBranch **bvp = new TBranch*[NFILESLHE];
  
  for(int i = 0; i<NFILESLHE;i++) {
    t[i]->SetBranchAddress("vdirection",&vdirection,&bvdirection[i]);
    t[i]->SetBranchAddress("veta",&veta,&bveta[i]);
    t[i]->SetBranchAddress("vphi",&vphi,&bvphi[i]);
    t[i]->SetBranchAddress("vM",&vM,&bvM[i]);
    t[i]->SetBranchAddress("vpdgId",&vpdgId,&bvpdgId[i]);
    t[i]->SetBranchAddress("vpT",&vpT,&bvpT[i]);
    t[i]->SetBranchAddress("vp",&vp,&bvp[i]);
  }

  TBranch **btyp = new TBranch*[NFILESLHCO];
  TBranch **beta = new TBranch*[NFILESLHCO];
  TBranch **bphi = new TBranch*[NFILESLHCO];
  TBranch **bjmas = new TBranch*[NFILESLHCO];
  TBranch **bntrk = new TBranch*[NFILESLHCO];
  TBranch **bpt = new TBranch*[NFILESLHCO];
  TBranch **bbtag = new TBranch*[NFILESLHCO];
  TBranch **bhadEm = new TBranch*[NFILESLHCO];

  for(int i = 0; i<NFILESLHCO;i++) {
    t2[i]->SetBranchAddress("typ",&typ,&btyp[i]);
    t2[i]->SetBranchAddress("eta",&eta,&beta[i]);
    t2[i]->SetBranchAddress("phi",&phi,&bphi[i]);
    t2[i]->SetBranchAddress("jmas",&jmas,&bjmas[i]);
    t2[i]->SetBranchAddress("ntrk",&ntrk,&bntrk[i]);
    t2[i]->SetBranchAddress("pt",&pt,&bpt[i]);
    t2[i]->SetBranchAddress("btag",&btag,&bbtag[i]);
    t2[i]->SetBranchAddress("hadEm",&hadEm,&bhadEm[i]);
  }

  

  const int NHISTS = 12;
  TH1F*** hists=new TH1F**[NHISTS];
  for (int ihist=0; ihist<NHISTS; ihist++) {
    hists[ihist]=new TH1F*[2];
  }
  
  const char * histTitles [NHISTS] = {"#eta of muons",
				      "#Delta#eta of muons",
				      "#phi of muons",
				      "pT of dimuon system",
				      "pdgId/types",
				      "pT of muons",
				      "#Delta#phi of muons",
				      "R of muons",
				      "m of dimuon system",
				      "p of dimuon system",
                                      "pT of leading muons",
				      "pT of subleading muons"};
  const char * yAxis = "Normalized counts/bin";
  const char *  axisLabels [NHISTS][2] = {{"#eta",yAxis},
					  {"#Delta#eta",yAxis},
					  {"#phi",yAxis},
					  {"pT (GeV/c)",yAxis},
					  {"pdgId",yAxis},
					  {"pT (GeV/c)",yAxis},
					  {"pT (GeV/c)",yAxis},
					  {"R",yAxis},
					  {"m (GeV/c^{2}}",yAxis},
					  {"p (GeV/c}",yAxis},
                                          {"pT (GeV/c)",yAxis},
					  {"pT (GeV/c)",yAxis}};
  const int defaultBins = 100;
  int logy  [NHISTS] = {0,
			0,
			0,
			0,
			1,
			0,
			0,
			0,
			0,
			0,
                        0,
                        0};
  int normalize  [NHISTS] = {1,
			     1,
			     1,
			     1,
			     1,
			     1,
			     1,
			     1,
			     1,
			     1,
			     1,
                             1};
  int nBins  [NHISTS] = {defaultBins,
			 defaultBins,
			 defaultBins,
			 defaultBins,
			 defaultBins,
			 defaultBins,
			 defaultBins,
			 defaultBins,
			 defaultBins,
			 defaultBins,
			 defaultBins,
			 defaultBins};
  float xLimits  [NHISTS] [2] = {{-7,7},
				 {-0.5,5.5},
				 {-0.5,6.5},
				 {0,150},
				 {-30,30},
				 {0,100},
				 {-2,7},
				 {-2,7},
				 {80,130},
				 {0,1000},
                                 {0,150},
				 {0,150}};
  float yLimit  [NHISTS];

  const char * histStart = "h";
  std::ostringstream ss;
  for(int i = 0; i < NHISTS; i++) {
    for(int j = 0; j < 2; j++) {
      ss.str("");
      ss << histStart << i+1 << "_" << j;
      hists[i][j] = new TH1F((ss.str()).c_str(),histTitles[i],nBins[i],xLimits[i][0],xLimits[i][1]);
    }

  }
  int nnonleptons, nLeptons = 0;
  float bosonMass, EFirstLepton, currentE,etaFirstLepton, currentEta, phiFirstLepton, currentPhi, dEta, dPhi = 0;
  float pTx, pTy, pz,p,pTFirstLepton, currentpT = 0;
  const float muonMass = 0.105658;
  int entriest = 0;
  int entriest2 = 0;
  for (int k = 0; k < NFILESLHE; k++) {
    entriest = (int)t[k]->GetEntries();
    for (int i = 0; i < entriest; i++) {
      nnonleptons = 0;
      nLeptons = 0;
      long long tentry = t[k]->LoadTree(i);
      bvdirection[k]->GetEntry(tentry);
      bveta[k]->GetEntry(tentry);
      bvM[k]->GetEntry(tentry);
      bvpdgId[k]->GetEntry(tentry);
      bvpT[k]->GetEntry(tentry);
      bvp[k]->GetEntry(tentry);
      bvphi[k]->GetEntry(tentry);
      EFirstLepton = 0.0;
      etaFirstLepton = 0.0;
      phiFirstLepton = 0.0;
      pTFirstLepton = 0.0;
      for (unsigned int j = 0; j < vdirection->size(); ++j) {
	if(abs(vpdgId->at(j))==13) {
	  hists[0][k]->Fill(veta->at(j));          
	  hists[5][k]->Fill(vpT->at(j));
	  hists[2][k]->Fill(vphi->at(j));
	  if(nLeptons==1) {
	    dEta = std::abs(etaFirstLepton-veta->at(j));
	    hists[1][k]->Fill(dEta);
	    if(vpT->at(j)>=pTFirstLepton) {
	      hists[10][k]->Fill(vpT->at(j));
	      hists[11][k]->Fill(pTFirstLepton);
	    }
	    else {
	      hists[10][k]->Fill(pTFirstLepton);
	      hists[11][k]->Fill(vpT->at(j));
	    }
	    pTx = std::cos(phiFirstLepton)*pTFirstLepton+std::cos(vphi->at(j))*vpT->at(j);
	    pTy = std::sin(phiFirstLepton)*pTFirstLepton+std::sin(vphi->at(j))*vpT->at(j);
	    pz = getPz(etaFirstLepton,pTFirstLepton)+getPz(veta->at(j),vpT->at(j));
	    float pz2 = std::sqrt(vp->at(j)*vp->at(j)-vpT->at(j)*vpT->at(j));
	    if(i<5) {
	      std::cout << "i: " << i << " pz1: "<< pz2 <<" pz2: " <<getPz(veta->at(j),vpT->at(j)) << std::endl;
	    }
	    hists[3][k]->Fill(std::sqrt(pTx*pTx+pTy*pTy));     
	    dPhi = std::abs(phiFirstLepton-vphi->at(j));
	    hists[6][k]->Fill(dPhi);
	    hists[7][k]->Fill(std::sqrt(dEta*dEta+dPhi*dPhi));
	    
	    
	    
	    p = std::sqrt(pTx*pTx+pTy*pTy+pz*pz);
	    currentE = std::sqrt(vp->at(j)*vp->at(j)+muonMass*muonMass);
	    hists[8][k]->Fill(std::sqrt((currentE+EFirstLepton)*(currentE+EFirstLepton)-p*p));
	    hists[9][k]->Fill(p);
	  }
	  etaFirstLepton = veta->at(j);
	  phiFirstLepton = vphi->at(j);
	  pTFirstLepton = vpT->at(j);
	  EFirstLepton = std::sqrt(vp->at(j)*vp->at(j)+muonMass*muonMass);
	 
	}
	if(vdirection->at(j)==-1) {
	  hists[4][k]->Fill(vpdgId->at(j));  
	} 
	if(abs(vpdgId->at(j))!=13) {
	  nnonleptons++;
	}
	else {
	  nLeptons++;
	}
      }
     
    }
  }
  
  for (int k = 0; k < NFILESLHCO; k++) {
    entriest2 = (int)t2[k]->GetEntries();
    for (int i = 0; i < entriest2; i++) {
      long long tentry2 = t2[k]->LoadTree(i);
      btyp[k]->GetEntry(tentry2);
      beta[k]->GetEntry(tentry2);
      bphi[k]->GetEntry(tentry2);
      bpt[k]->GetEntry(tentry2);
      bjmas[k]->GetEntry(tentry2);
      bntrk[k]->GetEntry(tentry2);
      bbtag[k]->GetEntry(tentry2);
      bhadEm[k]->GetEntry(tentry2);
      bosonMass = 0.0;
      nLeptons = 0;
      etaFirstLepton = 0.0;
      currentEta = 0.0;
      pTFirstLepton = 0.0;
      currentpT = 0.0;
      EFirstLepton = 0.0;
      currentE = 0.0;
      phiFirstLepton = 0.0;
      currentPhi = 0.0;
      for (unsigned int j = 0; j < typ->size(); ++j) {
	if((typ->at(j))==2) {
	  hists[5][k+NFILESLHE]->Fill(pt->at(j));
	  hists[0][k+NFILESLHE]->Fill(eta->at(j));
	  hists[2][k+NFILESLHE]->Fill(phi->at(j));
	  bosonMass += pt->at(j);
	  nLeptons++;
	  etaFirstLepton = currentEta;
	  currentEta = eta->at(j);
	  pTFirstLepton = currentpT;
	  EFirstLepton = currentE;

	  pz = getPz(eta->at(j), pt->at(j));
	  p = std::sqrt(pz*pz+pt->at(j)*pt->at(j));
	  currentE = std::sqrt(p*p+muonMass*muonMass);

	  currentpT = pt->at(j);
	  phiFirstLepton = currentPhi;
	  currentPhi = phi->at(j);
	}
	hists[4][k+NFILESLHE]->Fill(typ->at(j));
      }
      if(nLeptons==2) {
	if(currentpT>=pTFirstLepton) {
	  hists[10][k+NFILESLHE]->Fill(currentpT);
	  hists[11][k+NFILESLHE]->Fill(pTFirstLepton);
	}
	else {
	  hists[10][k+NFILESLHE]->Fill(pTFirstLepton);
	  hists[11][k+NFILESLHE]->Fill(currentpT);
	}
	pTx = std::cos(phiFirstLepton)*pTFirstLepton+std::cos(currentPhi)*currentpT;
	pTy = std::sin(phiFirstLepton)*pTFirstLepton+std::sin(currentPhi)*currentpT;
	pz = getPz(etaFirstLepton,pTFirstLepton)+ getPz(currentEta,currentpT);
	dEta = std::abs(etaFirstLepton-currentEta);
	hists[1][k+NFILESLHE]->Fill(dEta);
	hists[3][k+NFILESLHE]->Fill(std::sqrt(pTx*pTx+pTy*pTy)); 
	
	
	dPhi = std::abs(phiFirstLepton-currentPhi);
	hists[6][k+NFILESLHE]->Fill(dPhi);
	hists[7][k+NFILESLHE]->Fill(std::sqrt(dEta*dEta+dPhi*dPhi));

	
	p = std::sqrt(pTx*pTx+pTy*pTy+pz*pz);
	hists[8][k+NFILESLHE]->Fill(std::sqrt((currentE+EFirstLepton)*(currentE+EFirstLepton)-p*p));
	hists[9][k+NFILESLHE]->Fill(p);

      }
    }
  }
  
  for(int i = 0; i < NHISTS; i++) {
    if(normalize[i]==1) {
      for(int j = 0; j<NFILES;j++) {
	hists[i][j]->Scale(1.0/hists[i][j]->Integral());
      }
    }
  }

  double maximum = 0.0;
  double lmax = 0.0;
  for(int i = 0; i < NHISTS; i++) {
    maximum = 0.0;
    for(int j = 0; j<NFILES;j++) {
      hists[i][j]->SetXTitle(axisLabels[i][0]);
      hists[i][j]->SetYTitle(axisLabels[i][1]);
      lmax = hists[i][j]->GetMaximum();
      if(lmax>maximum) {
	maximum = lmax;
      }
    }
    yLimit[i]=maximum;
  }

  const char* fileExtension = ".pdf";
  const char* fileStart = "outputCombined/hist";
  std::ostringstream oss;

  myStyle->cd();
  myStyle->SetOptStat(0);
  TLegend* leg = new TLegend(0.7,0.7,0.85,0.85);
  leg->SetFillColor(0);

  /*----------------LEGEND--------------------*/
  leg->SetHeader(descriptions[2]);
  leg->AddEntry(hists[0][0],descriptions[0],"l");
  leg->AddEntry(hists[0][1],descriptions[1],"l");
  /*------------------------------------------*/

 
  for(int i = 0; i < NHISTS; i++) {
    hists[i][0]->SetMaximum(yLimit[i]*1.2);
    hists[i][1]->SetMaximum(yLimit[i]*1.2);
    hists[i][0]->SetLineColor(4);
    hists[i][0]->Draw();
    hists[i][1]->SetLineColor(2);
    hists[i][1]->Draw("same"); //"same"
    leg->Draw();
    c1->SetLogy(logy[i]);
    c1->Modified();
    c1->Update();
    oss.str("");
    oss << fileStart << i << fileExtension;
    c1->SaveAs((oss.str()).c_str(),"");

  }

  for(int i = 0; i<NFILESLHE;i++) {
    t[i]->ResetBranchAddresses();

  }
  for(int i = 0; i<NFILESLHCO;i++) {
    t2[i]->ResetBranchAddresses();

  }


  /* hists file */
  TFile *histsFile = TFile::Open("hists.root","RECREATE");
  for(int i = 0; i < NHISTS; i++) {
    hists[i][0]->Write();
    hists[i][1]->Write();
  }
  histsFile->Close(); 
}

int main(int argc, char **argv) {
  std::cout << "Program started!" << std::endl;
  histogramCreationCombined(&argv[1],&argv[3],&argv[5]);
  std::cout << "Program finished!" << std::endl;
}
